package org.travel_system.controller;

import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.Semaphore;

import jakarta.annotation.PostConstruct;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.ResultCallback;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.utils.JsonUtils;
import io.reactivex.Flowable;
import com.alibaba.dashscope.utils.Constants;
import org.travel_system.pojo.UserMessage;

@SpringBootApplication
@RestController
@RequestMapping("/ai")
public class Main {

    private static final Logger logger = LoggerFactory.getLogger(Main.class);

    @Value("${dashscope.api.key}")
    private String apiKey;

    public static void main(String[] args) {
        SpringApplication.run(Main.class, args);
    }

    @PostConstruct
    public void init() {
        // 设置 API 密钥
        Constants.apiKey = apiKey;
    }

    @PostMapping("/generate")
    public ResponseEntity<?> generate(@RequestBody UserMessage userMessage) {
        try {
            Generation gen = new Generation();
            Message userMsg = Message.builder().role(Role.USER.getValue()).content(userMessage.getContent()).build();

            String generatedContent = streamCallWithMessage(gen, userMsg);

            // 将生成的内容放入Map对象中返回
            Map<String, String> response = new HashMap<>();
            response.put("content", generatedContent);
            return ResponseEntity.ok().body(response);
        } catch (ApiException e) {
            logger.error("发生异常: {}", e.getMessage());
            return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR).body("生成出错: " + e.getMessage());
        } catch (Exception e) {
            logger.error("发生意外错误: {}", e.getMessage());
            return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR).body("发生意外错误.");
        }
    }

    private static String streamCallWithMessage(Generation gen, Message userMsg) {
        try {
            GenerationParam param = buildGenerationParam(userMsg);
            Flowable<GenerationResult> result = gen.streamCall(param);
            StringBuilder fullContent = new StringBuilder();

            result.blockingForEach(message -> handleGenerationResult(message, fullContent));

            logger.info("Full content: \n{}", fullContent.toString());
            return fullContent.toString();
        } catch (NoApiKeyException | ApiException | InputRequiredException e) {
            logger.error("An exception occurred during streamCall: {}", e.getMessage());
            return "An error occurred during generation: " + e.getMessage();
        } catch (Exception e) {
            logger.error("An unexpected error occurred during streamCall: {}", e.getMessage());
            return "An unexpected error occurred during generation.";
        }
    }


    private static void handleGenerationResult(GenerationResult message, StringBuilder fullContent) {
        fullContent.append(message.getOutput().getChoices().get(0).getMessage().getContent());
        logger.info("Received message: {}", JsonUtils.toJson(message));
    }

    private static GenerationParam buildGenerationParam(Message userMsg) {
        return GenerationParam.builder()
                .model("qwen-turbo")
                .messages(Arrays.asList(userMsg))
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .topP(0.8)
                .incrementalOutput(true)
                .build();
    }
}
